import cv2
import os
import numpy as np

# 定义类别名称列表，需要根据实际数据集修改
classes = ['pig', 'person']

# 数据集图像文件夹路径
image_folder = r'C:\Users\kang_\Desktop\d'
# 数据集标注文件夹路径
label_folder = r'C:\Users\kang_\Desktop\d'

# 获取图像文件列表
image_files = [f for f in os.listdir(image_folder) if f.endswith(('.jpg', '.png'))]

for image_file in image_files:
    # 构建图像文件的完整路径
    image_path = os.path.join(image_folder, image_file)
    # 读取图像
    image = cv2.imread(image_path)
    # 获取图像的高度和宽度
    height, width, _ = image.shape

    # 构建对应的标注文件路径
    label_file = os.path.splitext(image_file)[0] + '.txt'
    label_path = os.path.join(label_folder, label_file)
    if os.path.exists(label_path):
        # 打开标注文件
        with open(label_path, 'r') as f:
            lines = f.readlines()
            for line in lines:
                if len(line) < 2:
                    continue
                # 解析标注信息
                parts = line.strip().split()
                class_id = int(parts[0])
                # 分割掩码的归一化坐标
                seg_points = [float(p) for p in parts[1:]]
                # 将归一化坐标转换为像素坐标
                seg_points_pixel = []
                for i in range(0, len(seg_points), 2):
                    x = int(seg_points[i] * width)
                    y = int(seg_points[i + 1] * height)
                    seg_points_pixel.extend([x, y])
                # 将像素坐标转换为NumPy数组
                seg_points_pixel = np.array(seg_points_pixel, np.int32).reshape((-1, 1, 2))
                # 在图像上绘制分割掩码
                color = [(0, 255, 0), (255, 0, 0)]
                cv2.polylines(image, [seg_points_pixel], isClosed=True, color=color[class_id], thickness=2)
                # 在分割掩码上方绘制类别名称
                className = str(class_id)
                if class_id < len(classes):
                    className = classes[class_id]
                cv2.putText(image, className, (seg_points_pixel[0][0][0], seg_points_pixel[0][0][1] - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)

    # 显示图像
    cv2.imshow('YOLO Segment Dataset Preview', image)
    # 按任意键显示下一张图像
    cv2.waitKey(0)

# 关闭所有OpenCV窗口
cv2.destroyAllWindows()
